# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：grid_generator.py
    Created on：2025/6/7 15:10
    Author：(input)
    Description:
"""
import numpy as np

class BaseGrid:
    """
    网格基类，定义标准接口与网格数据结构。
    """
    def __init__(self):
        self.X = None  # 实部坐标网格
        self.Y = None  # 虚部坐标网格
        self.Z = None  # 复平面网格

    def get_real(self):
        return self.X

    def get_imag(self):
        return self.Y

    def get_complex(self):
        return self.Z

    def to_polar(self):
        r = np.abs(self.Z)
        theta = np.angle(self.Z)
        return r, theta


class RectangularGrid(BaseGrid):
    """
    规则笛卡尔网格类。
    """
    def __init__(self, x_length, y_length, Nx, Ny=None, origin=(0.0, 0.0)):
        super().__init__()
        self.x_length = x_length
        self.y_length = y_length
        self.Nx = Nx
        self.Ny = Ny if Ny else Nx
        self.origin = origin

        self.generate_grid()

    def generate_grid(self):
        x_start = -self.origin[0]
        x_end = self.x_length - self.origin[0]
        y_start = -self.origin[1]
        y_end = self.y_length - self.origin[1]

        x = np.linspace(x_start, x_end, self.Nx)
        y = np.linspace(y_start, y_end, self.Ny)
        self.X, self.Y = np.meshgrid(x, y)
        self.Z = self.X + 1j * self.Y


class MappedGrid(BaseGrid):
    """
    映射变换后的规则网格。
    由一个已有网格 + 映射函数（如Möbius、双线性）生成变换后的网格。
    """
    def __init__(self, base_grid: BaseGrid, map_func):
        super().__init__()
        self.base_grid = base_grid
        self.map_func = map_func
        self.apply_mapping()

    def apply_mapping(self):
        Z0 = self.base_grid.get_complex()
        W = self.map_func(Z0)  # 支持 vectorized 函数
        self.Z = W
        self.X = W.real
        self.Y = W.imag


class BodyFittedGrid(BaseGrid):
    """
    非规则贴体网格类，使用双线性插值构造。
    """
    def __init__(self, corner_points, m=50, n=50):
        """
        corner_points: list of 4 points in order: [A1, A2, A3, A4]
            A1: top-right, A2: top-left, A3: bottom-left, A4: bottom-right
        """
        super().__init__()
        self.P = np.array(corner_points)  # 4个顶点 (x, y)
        self.m = m
        self.n = n
        self.generate_grid()

    def shape_functions(self, xi, eta):
        N0 = xi * eta      # A1 (1,1)
        N1 = (1 - xi) * eta    # A2 (0,1)
        N2 = (1 - xi) * (1 - eta)  # A3 (0,0)
        N3 = xi * (1 - eta)   # A4 (1,0)
        return np.array([N0, N1, N2, N3])

    def generate_grid(self):
        xi_vals = np.linspace(0, 1, self.m)
        eta_vals = np.linspace(0, 1, self.n)
        Xi, Eta = np.meshgrid(xi_vals, eta_vals, indexing='ij')

        self.X = np.zeros((self.m, self.n))
        self.Y = np.zeros((self.m, self.n))

        for i in range(self.m):
            for j in range(self.n):
                N = self.shape_functions(Xi[i, j], Eta[i, j])
                x = np.dot(N, self.P[:, 0])
                y = np.dot(N, self.P[:, 1])
                self.X[i, j] = x
                self.Y[i, j] = y

        self.Z = self.X + 1j * self.Y
